🚀 Deployed on https://68bb3aa8cd486a0a579d9ae4--nostalgic-wright-5fa355.netlify.app

📖 Netlify Preview Ready!

Preview URL: https://pr-591--sunny-cactus-210e3e.netlify.app (5b0ce92)

📚 Changed Lecture Pages: kesten_processes

@kp992 @HumphreyYang I haven't come across this error before

Unable to initialize backend 'cuda': INVALID_ARGUMENT: Unexpected option name passed to PJRT_Client_Create: use_tfrt_gpu_client (you may need to uninstall the failing plugin package, or set JAX_PLATFORMS=cpu to skip this backend

I don't see this happening in other PRs -- but looks like an install issue?

It is coming from

@jit
def generate_draws(
    key=random.PRNGKey(123),
    μ_a=-0.5,
    σ_a=0.1,
    μ_b=0.0,
    σ_b=0.5,
    μ_e=0.0,
    σ_e=0.5,
    s_bar=1.0,
    T=500,
    M=1_000_000,
    s_init=1.0,
):
    """
    JAX-jit version of the generate_draws function.
    Returns:
        Array of M draws
    """
    # Create M different random keys for parallel execution
    keys = random.split(key, M)

    draws = vectorized_single_draw(
        keys, μ_a, σ_a, μ_b, σ_b, μ_e, σ_e, s_bar, T, s_init
    )

    return draws

Hi @mmcky,

The code runs well on my GPU, I think this might be related to the environment. I will retrigger the action and see if it helps.

📖 Netlify Preview Ready!

Preview URL: https://pr-591--sunny-cactus-210e3e.netlify.app (6e44c3c)

📚 Changed Lecture Pages: kesten_processes

Hi @mmcky,

This is fixed now. It's a bit tricky to find but it is caused by a minor typo :tags: [no-execute] -> :tags: [skip-execution].

The code runs well on my GPU, I think this might be related to the environment.

Thanks for fixing this. Even for me it worked fine. Probably we can work on making debug messages easier.

The code runs well on my GPU, I think this might be related to the environment.

Thanks for fixing this. Even for me it worked fine. Probably we can work on making debug messages easier.

Many thanks @kp992, I was expecting bug like this to be raised by jupyter but not jax. It's interesting it is raised as a device interface issue : )

The code runs well on my GPU, I think this might be related to the environment.

Thanks for fixing this. Even for me it worked fine. Probably we can work on making debug messages easier.

Many thanks @kp992, I was expecting bug like this to be raised by jupyter but not jax. It's interesting it is raised as a device interface issue : )

The no-execute mean that that cell was still executed -- as it is an invalid tag. Nice work @HumphreyYang

thanks @kp992 you read my mind! Was just about to make the edit :)

Will merge this soon. Thanks for your PR.

@kp992 once this finishes building -- can you please do a final output check of the preview link and give me your thumbs up.

📖 Netlify Preview Ready!

Preview URL: https://pr-591--sunny-cactus-210e3e.netlify.app (1adfc93)

📚 Changed Lecture Pages: kesten_processes

Looks good to me @mmcky

📖 Netlify Preview Ready!

Preview URL: https://pr-591--sunny-cactus-210e3e.netlify.app (9bead7e)

📚 Changed Lecture Pages: kesten_processes

@mmcky I'm not seeing any JAX on the 'Changed lecture pages' link above.

@mmcky I'm not seeing any JAX on the 'Changed lecture pages' link above.

@jstac the jax code in this lecture is contained in the exercise solutions - but do let me know if you don't see it on your end.

We have imports in the solution

We typically have a rule that all imports are at the top of the lecture, however I am in two minds re: solutions. Half of me thinks keeping imports allowed in solutions can be useful when introducing a new package that hasn't been used in the lecture itself. Any thoughts on this style-guide item @jstac

Ah, I see. Thanks @mmcky

We should shift this lecture to JAX throughout.

All simulation of paths in JAX should use the generate_paths functionality from the lake model PR #589

Many thanks @kp992, @jstac, and @mmcky!

I will have a go at converting them into the new template!

thanks @HumphreyYang I have started on the documentation in the manual using the previous PR.

I will use this one as a final cross check and then do a final edit of https://github.com/QuantEcon/QuantEcon.manual/pull/73/files.

📖 Netlify Preview Ready!

Preview URL: https://pr-591--sunny-cactus-210e3e.netlify.app (7d00832)

📚 Changed Lecture Pages: kesten_processes

thanks @HumphreyYang for these changes.

@jstac -- I am in favour of consistency, but this change seems to be adding quite a few lines of code (making it more complex) -- so I just wanted to run this template by you to make sure we are heading down the right path here.

@mmcky This is, in a sense, the cost of using JAX. JAX wasn't built for iteration and has no immutable arrays.

If we're going to keep using the lax.scan logic, then, given it's slightly unusual syntax, I think it's better to pull that logic out into a specialized function and explain it.

I recommend not hiding this code, since it's essential to the lecture.

I'm very open to other ideas and approaches.

The other option is refrain from JAX conversions for lectures that iterate a low dimensional object forward many times. Admittedly, numpy and numba are better for this. I'm open to this discussion too.

Thanks @jstac.

If we're going to keep using the lax.scan logic, then, given it's slightly unusual syntax, I think it's better to pull that logic out into a specialized function and explain it.

Completely agree.

I recommend not hiding this code, since it's essential to the lecture.

👍

I'm very open to other ideas and approaches.

I guess I would like to have tabbed code blocks available - giving the choice (But that is more to maintain, so committing is probably optimal)

The other option is refrain from JAX conversions for lectures that iterate a low dimensional object forward many times. Admittedly, numpy and numba are better for this. I'm open to this discussion too.

Are you thinking of something like if lecture > 10s runtime then use jax otherwise use numba kind of rule to figure out when to apply it?

I will mull over this tonight.

@HumphreyYang you've done a great job with this conversion. It is going to be a really useful pattern to finalise the docs.

📖 Netlify Preview Ready!

Preview URL: https://pr-591--sunny-cactus-210e3e.netlify.app (4f4c0f5)

📚 Changed Lecture Pages: kesten_processes